// This is a small program to precompile hlsl library to a header file,
// used only for fixed function shader generation with d3d11 linker. 
// It also generates 8 * 4 combinations of texture sampling functions.
//
// Useful because fxc doesn't support lib_* targets,
// but once it does it's probably a good idea to deprecate this.

#include <windows.h>
#include <d3d11.h>
#include <d3dcommon.h>
#include <d3dcompiler.h>

#include <cstdio>
#include <cassert>

typedef HRESULT (WINAPI *D3DCompileFunc)(
	const void* pSrcData,
	unsigned long SrcDataSize,
	const char* pFileName,
	const D3D10_SHADER_MACRO* pDefines,
	void* pInclude,
	const char* pEntrypoint,
	const char* pTarget,
	unsigned int Flags1,
	unsigned int Flags2,
	ID3DBlob** ppCode,
	ID3DBlob** ppErrorMsgs
);

const char* dllName = "D3DCompiler_47.dll";

const char* readFile(const char* path)
{
	FILE* f = fopen(path, "rb");
	if (!f)
		return NULL;

	fseek(f, 0, SEEK_END);
	size_t s = ftell(f);
	fseek(f, 0, SEEK_SET);
	char* out = (char*)malloc(s+1);
	fread(out, s, 1, f);
	out[s] = '\0';
	fclose(f);
	return out;
}

void printCompiledShader(ID3DBlob* shader, const char* name, FILE* out)
{
	const unsigned char* data = (const unsigned char*)shader->GetBufferPointer();
	size_t size = (size_t)shader->GetBufferSize();

	fprintf(out, "const BYTE %s[] = { \n", name);
	for(int i = 0; i < size; ++i) {
		fprintf(out, ((i==size-1) ? "%d" : "%d,"), (int)data[i]);
		if((i+1) % 16 == 0)
			fprintf(out, "\n");
	}
	fprintf(out, "\n};\n");
}

// type: 0 - 2d, 1 - 2d proj, 2 - 3d, 3 - cube
const char* generateTextureSamplingLib(int type, int unit)
{
	const int maxSize = 2048;
	static char hlsl[maxSize];

	const char* textureTypes[] = { "Texture2D", "Texture2D", "Texture3D", "TextureCube" };
	const char* coordTypes[] = {"float2", "float4", "float3", "float3" };
	const char* coords[] = {"uv.xy", "uv.xy/uv.w", "uv.xyz", "uv.xyz" };

	sprintf(hlsl,
		"%s<float4> Tex%d : register(t%d);\n"
		"SamplerState Smp%d : register(s%d);\n"
		"export float4 LoadTex%d(%s uv) { return Tex%d.Sample(Smp%d, %s); }\n",
		textureTypes[type], unit, unit, unit, unit, unit,
		coordTypes[type], unit, unit, coords[type]
	);

	return hlsl;
}

const char* textureSamplingLibName(int type, int unit)
{
	static const char* typeStr[] = {"2D", "Proj", "3D", "Cube"};
	static char name[32];
	sprintf(name, "g_FFSampleTex%s%d", typeStr[type], unit);
	return name;
}

// Since we can't bind multiple textures to the same texture register,
// we'll have to generate 4 * 8 sampling functions and link the correct ones
// at runtime.
void printTextureSampling(D3DCompileFunc compileFunc, FILE* out)
{
	// Print compiled shaders
	for (int type = 0; type < 4; ++type)
	{
		for (int unit = 0; unit < 8; ++unit)
		{
			const char* name = textureSamplingLibName(type, unit);
			const char* hlsl = generateTextureSamplingLib(type, unit);

			ID3DBlob* errors;
			ID3DBlob* shader;
			HRESULT hr = compileFunc(
				hlsl, strlen(hlsl), name, NULL, NULL, NULL,
				"lib_4_0_level_9_1", D3DCOMPILE_OPTIMIZATION_LEVEL3, 0, &shader, &errors
			);

			assert(SUCCEEDED(hr));

			printCompiledShader(shader, name, out);

			shader->Release();
		}
	}

	// Print index to all shaders
	fprintf(out, "const BYTE* g_FFSampleTexLib[] = { \n");
	for (int type = 0; type < 4; ++type)
		for (int unit = 0; unit < 8; ++unit)
			fprintf(out, "%s,\n", textureSamplingLibName(type, unit));
	fprintf(out, "};\n");

	// Print sizes of all shaders
	fprintf(out, "const size_t g_FFSampleTexLibSize[] = { \n");
	for (int type = 0; type < 4; ++type)
		for (int unit = 0; unit < 8; ++unit)
			fprintf(out, "sizeof(%s),\n", textureSamplingLibName(type, unit));
	fprintf(out, "};\n");
}

int main(int argc, const char** argv)
{
	if (argc != 3)
	{
		printf("Usage: CompileShaderLib.exe src.hlsl out.h\n");
		return 1;
	}

	HMODULE dll = LoadLibraryA (dllName);
	if (!dll)
	{
		printf("Can't load %s\n", dllName);
		return 1;
	}

	D3DCompileFunc compileFunc = (D3DCompileFunc) GetProcAddress(dll, "D3DCompile");
	if (!compileFunc)
	{
		printf("Can't get D3DCompile function address\n");
		return 1;
	}

	const char* hlsl = readFile(argv[1]);

	ID3DBlob* errors;
	ID3DBlob* shader;
	HRESULT hr = compileFunc(
		hlsl, strlen(hlsl), argv[1], NULL, NULL, NULL,
		"lib_4_0_level_9_1", D3DCOMPILE_OPTIMIZATION_LEVEL3, 0, &shader, &errors
	);

	if (FAILED(hr))
	{
		printf("Failed to compile, 0x%x:\n%s\n", hr, errors ? errors->GetBufferPointer() : "");
		return 1;
	}
	
	free((void*)hlsl);
	FILE* out = fopen(argv[2], "w");
	
	fprintf(out, "// File autogenerated by CompileShaderLib.exe\n");
	
	printCompiledShader(shader, "g_FFShaderLibrary", out);
	shader->Release();

	printTextureSampling(compileFunc, out);

	fclose(out);

	return 0;
}

